/******************************************************************************
 * Copyright (c) 2022 Jaroslav Hensl                                          *
 *                                                                            *
 * Permission is hereby granted, free of charge, to any person                *
 * obtaining a copy of this software and associated documentation             *
 * files (the "Software"), to deal in the Software without                    *
 * restriction, including without limitation the rights to use,               *
 * copy, modify, merge, publish, distribute, sublicense, and/or sell          *
 * copies of the Software, and to permit persons to whom the                  *
 * Software is furnished to do so, subject to the following                   *
 * conditions:                                                                *
 *                                                                            *
 * The above copyright notice and this permission notice shall be             *
 * included in all copies or substantial portions of the Software.            *
 *                                                                            *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,            *
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES            *
 * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND                   *
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT                *
 * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,               *
 * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING               *
 * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR              *
 * OTHER DEALINGS IN THE SOFTWARE.                                            *
 *                                                                            *
*******************************************************************************/
#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>
#include <bitstream.h>
#include <bpatcher.h>

/* data */
#include "fasmdiff.h"
#include "fasmdiff_v2.h"
#include "fasmdiff_me.h"

#include "../cpuspeed/speed_v1_diff.h"
#include "../cpuspeed/speed_v2_diff.h"
#include "../cpuspeed/speed_v3_diff.h"
#include "../cpuspeed/speed_v4_diff.h"

#include "../cpuspeed/speed_v5_diff.h"
#include "../cpuspeed/speed_v6_diff.h"
#include "../cpuspeed/speed_v7_diff.h"
#include "../cpuspeed/speed_v8_diff.h"

#include "../cpuspeed/speedndis_v1_diff.h"
#include "../cpuspeed/speedndis_v2_diff.h"
#include "../cpuspeed/speedndis_v3_diff.h"
#include "../cpuspeed/speedndis_v4_diff.h"

#include "fasmdiff_old.h"
#include "fasmdiff_old_v2.h"

#define PREFIX_MAX 255

const char *h_header = \
	"/****** Generated by " __FILE__ " ******/\n" \
	"/******    changes may be lost   ******/\n\n" \
	"#ifndef __%s_H__INCLUDED__\n" \
	"#define __%s_H__INCLUDED__\n\n" \
	"#include <patch.h>\n\n";

const char *h_footer = "#endif /* __%s_H__INCLUDED__ */\n";

typedef struct _preset_t
{
	const char *name;
	const char *dumped;
	const char *original;
	const char *patched;
	const char *relocated;
	const uint8_t *dumped_diff;
	size_t dumped_diff_size;
} preset_t;

#define PRESENT(_dir, _name) {#_name, \
	_dir "/" #_name "_dump.bin", \
	_dir "/" #_name "_orig.bin", \
	_dir "/" #_name ".bin", \
  _dir "/" #_name "_reloc.bin", \
  vmm_ ##_name## _diff, sizeof(vmm_ ##_name## _diff)}

preset_t presets[] = {
	{"98",   "vmm/dump.bin",    "vmm/original.bin",    "vmm/patched.bin",    "vmm/reloc.bin",    vmm_fasmdiff,    sizeof(vmm_fasmdiff)},
	{"98v2", "vmm/dump_v2.bin", "vmm/original_v2.bin", "vmm/patched_v2.bin", "vmm/reloc_v2.bin", vmm_fasmdiff_v2, sizeof(vmm_fasmdiff_v2)},
	{"me",   "vmm/dumpme.bin",  "vmm/originalme.bin",  "vmm/patchedme.bin",  "vmm/relocme.bin",  vmm_fasmdiff_me, sizeof(vmm_fasmdiff_me)},
	
	{"98_old",    "vmm/dump_old.bin",       "vmm/original_old.bin",      "vmm/patched.bin",       "vmm/reloc_old.bin",     vmm_fasmdiff_old,    sizeof(vmm_fasmdiff_old)},
	{"98_old_v2", "vmm/dump_old_v2.bin",    "vmm/original_old_v2.bin",   "vmm/patched_v2.bin",    "vmm/reloc_old_v2.bin",  vmm_fasmdiff_old_v2, sizeof(vmm_fasmdiff_old_v2)},
	
	{"98_simple",    "vmm/dump.bin",    "vmm/original_simple.bin",    "vmm/patched_simple.bin",    "vmm/reloc_simple.bin",    vmm_fasmdiff,    sizeof(vmm_fasmdiff)},
	{"98_simple_v2", "vmm/dump_v2.bin", "vmm/original_simple_v2.bin", "vmm/patched_simple_v2.bin", "vmm/reloc_simple_v2.bin", vmm_fasmdiff_v2, sizeof(vmm_fasmdiff_v2)},

	PRESENT("cpuspeed", speed_v1),
	PRESENT("cpuspeed", speed_v2),
	PRESENT("cpuspeed", speed_v3),
	PRESENT("cpuspeed", speed_v4),
	
	PRESENT("cpuspeed", speed_v5),
	PRESENT("cpuspeed", speed_v6),
	PRESENT("cpuspeed", speed_v7),
	PRESENT("cpuspeed", speed_v8),
	
	PRESENT("cpuspeed", speedndis_v1),
	PRESENT("cpuspeed", speedndis_v2),
	PRESENT("cpuspeed", speedndis_v3),
	PRESENT("cpuspeed", speedndis_v4),

	{NULL, NULL, NULL, NULL, NULL, NULL, 0}
};

static void upper_case(const char *src, char *dst, size_t dstsize)
{
	size_t i;
	size_t len = strlen(src);
	if(len >= dstsize)
	{
		len = dstsize-1;
	}
	
	for(i = 0; i < len; i++)
	{
		dst[i] = toupper(src[i]);
	}
	dst[len] = '\0';
}

static size_t file_size(FILE *fp)
{
	size_t cur, end;
	
	cur = ftell(fp);
	fseek(fp, 0, SEEK_END);
	end = ftell(fp);
	fseek(fp, cur, SEEK_SET);
	
	return end;
}

int main(int argc, char *argv[])
{
	uint8_t *mem1 = NULL, *mem2 = NULL, *mem3 = NULL; /* buffer for fread */
	
	size_t dumped_s, original_s, patched_s, relocated_s, bin_size, patch_ofset, patch_size;
	
	FILE *dumped      = NULL; /* memory dump or copy from file original windows VMM.VXD */
	FILE *original    = NULL; /* assembled code with original sequence */
	FILE *patched     = NULL; /* assembled new code */
	FILE *relocated   = NULL; /* assembled original code but with rellocated all absolute addreseses */
	FILE *out         = NULL; /* out */
	
	void *patch_diff_mem = NULL, *reloc_diff_mem = NULL, *orig_diff_mem = NULL, *cmp_mem = NULL;
	
	bitstream_t patch_diff;   /* difference: original<->patched */
	bitstream_t reloc_diff; /* difference: original<->relocated */
	bitstream_t orig_diff;    /* difference: original<->dumped */
	bitstream_t cmp;          /* compare results */
	
	const char *prefix = NULL;
	char prefix_upper[PREFIX_MAX];
	
	preset_t *preset = NULL;
	
	int result = EXIT_FAILURE;

	if(argc < 7)
	{
		fprintf(stderr, "Usage: %s <preset> <prefix> <required size of binaries> <patch offset> <patch size> <output.h>\n", argv[0]);
		return EXIT_FAILURE;
	}
	else
	{
		prefix      = argv[2];
		upper_case(prefix, prefix_upper, PREFIX_MAX);
		bin_size    = strtoul(argv[3], NULL, 0);
		patch_ofset = strtoul(argv[4], NULL, 0);
		patch_size  = strtoul(argv[5], NULL, 0);
		
		preset = &presets[0];
		
		do
		{
			if(strcmp(argv[1], preset->name) == 0)
			{
				break;
			}
			
			preset++;
		} while(preset->name != NULL);
		
		if(preset == NULL || preset->name == NULL)
		{
			fprintf(stderr, "Preset %s in not present\n", argv[1]);
			return EXIT_FAILURE;
		}
	}

	dumped = fopen(preset->dumped, "rb");
	original = fopen(preset->original, "rb");
	patched = fopen(preset->patched, "rb");
	relocated = fopen(preset->relocated, "rb");
	out = fopen(argv[6], "wb");
	if(out == NULL)
	{
		fprintf(stderr, "Cannot open output %s\n", argv[6]);
		return EXIT_FAILURE;
	}
	
	if(original != NULL && patched != NULL && relocated != NULL)
	{
		dumped_s = 0;
		if(dumped)
		{
			dumped_s = file_size(dumped);
		}
		
		original_s   = file_size(original);
		patched_s    = file_size(patched);
		relocated_s = file_size(relocated);
		
		if(dumped != NULL && dumped_s != bin_size)
		{
			fprintf(stderr, "%s: invalid size\n", preset->dumped);
		}
		else if(original_s != bin_size)
		{
			fprintf(stderr, "%s: invalid size\n", preset->original);
		}
		else if(patched_s != bin_size)
		{
			fprintf(stderr, "%s: invalid size\n", preset->patched);
		}
		else if(relocated_s != bin_size)
		{
			fprintf(stderr, "%s: invalid size\n", preset->relocated);
		}
		else
		{
			size_t site_bs = bs_calc_size(bin_size);
			uint8_t buf8 = 0;
			size_t bin_size_round = (bin_size + 7) & 0xFFFFFFF8;
			
			/* buffers */
			mem1 = calloc(1, bin_size_round);
			mem2 = calloc(1, bin_size_round);
			mem3 = calloc(1, bin_size_round);
			
			/* bitstreams */
			patch_diff_mem   = bs_mem_alloc(&patch_diff, site_bs);
			reloc_diff_mem = bs_mem_alloc(&reloc_diff, site_bs);
			orig_diff_mem    = bs_mem_alloc(&orig_diff, site_bs);
			cmp_mem          = bs_mem_alloc(&cmp, site_bs);
			
			if(mem1 != NULL && mem2 != NULL && mem3 != NULL && 
				patch_diff_mem != NULL && reloc_diff_mem != NULL && orig_diff_mem != NULL &&
				cmp_mem != NULL)
			{
				size_t i, j;
				int test_realoc = 1;
				
				fread(mem1, 1, bin_size, original);
				fread(mem2, 1, bin_size, relocated);
				diff_sieve(mem1, mem2, bin_size_round, &reloc_diff);
				
				fread(mem2, 1, bin_size, patched);
				diff_sieve(mem1, mem2, bin_size_round, &patch_diff);
				
				if(dumped != NULL)
				{
					fread(mem3, 1, bin_size, dumped);
					diff_sieve(mem1, mem3, bin_size_round, &orig_diff);
				}
				else
				{
					fprintf(stderr, "Warning: %s missing, some tests skipped\n", preset->dumped);
					
					bs_mem_free(&orig_diff);
					orig_diff_mem = NULL;
					bs_mem(&orig_diff, (uint8_t*)preset->dumped_diff, preset->dumped_diff_size);
				}
				
				/* file header */
				fprintf(out, h_header, prefix_upper, prefix_upper);
				
				bs_reset(&patch_diff);
				bs_reset(&reloc_diff);
				bs_reset(&cmp);
				bs_logic(BS_AND, &patch_diff, &reloc_diff, &cmp, bin_size_round);
				bs_reset(&cmp);
				test_realoc = bs_is_zero(&cmp, bin_size_round);
				if(test_realoc == 0)
				{
					size_t i = 0;
					fprintf(stderr, "Realocation issues (one bit -> one byte):\n");
					for(i = 0; i < bs_calc_size(bin_size_round); i++)
					{
						fprintf(stderr, "%02X ", ((uint8_t*)cmp_mem)[i]);
					}
					fprintf(stderr, "\n");
					
				}
				
				bs_reset(&orig_diff);
				bs_reset(&reloc_diff);
				bs_reset(&cmp);
				bs_logic(BS_OR, &orig_diff, &reloc_diff, &cmp, bin_size_round);
								
				// orig
				fprintf(out, "/* original data for search */\nconst uint8_t %s_orig[] = {", prefix);
				for(i = 0; i < patch_size; i++)
				{
					if(i % 16 == 0)
					{
						fprintf(out, "\n\t");
					}
					fprintf(out, "0x%02X, ", mem1[patch_ofset+i]);
				}
				fprintf(out, "\n};\n\n");
				
				// orig_check
				bs_reset(&cmp);
				i = patch_ofset % 8;
				j = 0;
				if(i != 0)
				{
					buf8 = bs_read_bit(&cmp, i);
				}
				fprintf(out, "/* bitmap of bytes to check */\nconst uint8_t %s_orig_check[] = {\n\t", prefix);
				for(; i < bin_size_round; i += 8)
				{
				  buf8 = ~bs_read_bit(&cmp, 8);
				  if(i >= patch_ofset && i < patch_ofset+patch_size)
				  {
						if((j++) % 16 == 0)
						{
							fprintf(out, "\n\t");
						}
						fprintf(out, "0x%02X, ", buf8);
					}
				}
				fprintf(out, "\n};\n\n");
				
				// patch
				fprintf(out, "/* new data to replace */\nconst uint8_t %s[] = {", prefix);
				for(i = 0; i < patch_size; i++)
				{
					if(i % 16 == 0)
					{
						fprintf(out, "\n\t");
					}
					fprintf(out, "0x%02X, ", mem2[patch_ofset+i]);
				}
				fprintf(out, "\n};\n\n");
				
				// patch_modif
				bs_reset(&patch_diff);
				i = patch_ofset % 8;
				j = 0;
			
				if(i != 0)
				{
					buf8 = bs_read_bit(&patch_diff, i);
				}
				fprintf(out, "/* bitmap of modify bytes */\nconst uint8_t %s_modif[] = {", prefix);
				for(; i < bin_size_round; i += 8)
				{
				  buf8 = bs_read_bit(&patch_diff, 8);
				  if(i >= patch_ofset && i < patch_ofset+patch_size)
				  {
						if((j++) % 16 == 0)
						{
							fprintf(out, "\n\t");
						}
						fprintf(out, "0x%02X, ", buf8);
					}
				}
				fprintf(out, "\n};\n\n");
				
				/* cpatch struct */
				fprintf(out,
					"const cpatch_t %s_cp = {\n\t%s, sizeof(%s),\n\t%s_orig, sizeof(%s_orig),\n\t"
					"%s_orig_check, sizeof(%s_orig_check),\n\t%s_modif, sizeof(%s_modif)\n};\n\n",
					prefix,
					prefix, prefix, prefix, prefix,
					prefix, prefix, prefix, prefix);
				
				/* file end */
				fprintf(out, h_footer, prefix_upper);
				
				if(test_realoc == 0)
				{
					fprintf(stderr, "Error: Code changes on relocate symbols!\n");
					//result = EXIT_SUCCESS; // !
				}
				else
				{
					result = EXIT_SUCCESS;
				}
				
			}
			else
			{
				fprintf(stderr, "Out of memory\n");
			}
		}
	}
	else
	{
		fprintf(stderr, "Required files are: %s, %s and %s\n", preset->patched, preset->original, preset->relocated);
	}
	
	fclose(out);
	
	if(dumped != NULL)      fclose(dumped);
	if(original != NULL)    fclose(original);
	if(patched != NULL)	    fclose(patched);
	if(relocated != NULL) fclose(relocated);
	
	if(mem1 != NULL) free(mem1);
	if(mem2 != NULL) free(mem2);
	if(mem3 != NULL) free(mem3);
	
	if(patch_diff_mem)   bs_mem_free(&patch_diff);
	if(reloc_diff_mem)   bs_mem_free(&reloc_diff);
	if(orig_diff_mem)    bs_mem_free(&orig_diff);
	if(cmp_mem)          bs_mem_free(&cmp);

	return result;
}
